import pymysql
from pymysql.cursors import DictCursor

class VideoDataMigrator:
    def __init__(self, db_config):
        """初始化数据库连接配置"""
        self.db_config = db_config
        self.connection = None
    
    def connect(self):
        """建立数据库连接"""
        try:
            self.connection = pymysql.connect(
                host=self.db_config['host'],
                user=self.db_config['user'],
                password=self.db_config['password'],
                database=self.db_config['database'],
                port=self.db_config.get('port', 3306),
                cursorclass=DictCursor
            )
            return True
        except Exception as e:
            print(f"数据库连接失败: {str(e)}")
            return False
    
    def get_qualified_records(self):
        """获取dddj_video表中content_type=2的记录"""
        if not self.connection:
            print("请先建立数据库连接")
            return []
            
        try:
            with self.connection.cursor() as cursor:
                sql = "SELECT * FROM dddj_video WHERE content_type = 2"
                cursor.execute(sql)
                return cursor.fetchall()
        except Exception as e:
            print(f"查询数据失败: {str(e)}")
            return []
    
    def migrate_records(self, records):
        """将记录迁移到video表"""
        if not self.connection or not records:
            print("缺少数据库连接或迁移数据")
            return 0
            
        success_count = 0
        
        try:
            with self.connection.cursor() as cursor:
                # 准备插入SQL
                insert_sql = """
                INSERT INTO video 
                (video_id, author_id, title, path, publish_time, likes, reposts)
                VALUES (%s, %s, %s, %s, %s, %s, %s)
                """
                
                # 处理每条记录
                for record in records:
                    try:
                        # 准备要插入的数据
                        data = (
                            record['id'],                  # dddj_video.id -> video.video_id
                            record['user_id'],             # dddj_video.user_id -> video.author_id
                            record['text_content'],        # dddj_video.text_content -> video.title
                            record['video_origin_url'],    # dddj_video.video_origin_url -> video.path
                            record['create_time'],         # dddj_video.create_time -> video.publish_time
                            record['love_count'],          # dddj_video.love_count -> video.likes
                            record['reply_count']          # dddj_video.reply_count -> video.reposts
                        )
                        
                        cursor.execute(insert_sql, data)
                        success_count += 1
                    except Exception as e:
                        continue
                
                # 提交事务
                self.connection.commit()
                return success_count
                
        except Exception as e:
            self.connection.rollback()
            print(f"迁移数据失败: {str(e)}")
            return 0
    
    def close(self):
        """关闭数据库连接"""
        if self.connection:
            self.connection.close()
            self.connection = None
    
    def run(self):
        """执行完整的迁移流程"""
        if not self.connect():
            return False
            
        try:
            print("开始查询符合条件的记录...")
            records = self.get_qualified_records()
            
            if records:
                print("开始迁移数据...")
                success = self.migrate_records(records)
                print(f"迁移完成，新增迁移 {success} 条记录，已有 {len(records) - success} 条记录")
                return True
            else:
                print("没有需要迁移的记录")
                return True
                
        finally:
            self.close()

# 使用示例
if __name__ == "__main__":
    # 数据库配置，请根据实际情况修改
    db_config = {
        'host': 'localhost',
        'port': 3306,
        'user': 'root',
        'password': '123456',
        'database': 'mydb'
    }
    
    migrator = VideoDataMigrator(db_config)
    migrator.run()


